from pydantic import BaseModel, Field
from langchain.agents import tool

from utils.utils import extract_data
import json
from config import web_json


class bar_chart_input(BaseModel):
    title: str = Field(..., description="生成柱状图组件的标题")
    desc: str = Field(
        ...,
        description="适合柱状图的描述段落（如：漏洞等级高危、中危、低危数量）",
    )


bar_chart_desc = """\
提取适合生成柱状图的文本段落

要求：
1. 文本必须包含至少2个可对比的类别及其对应数值
2. 数值需为明确数字（如100、25%、$300等）
"""


@tool(description=bar_chart_desc, args_schema=bar_chart_input)
def bar_chart(*args, **kwargs) -> str:
    title = ""
    if "title" in kwargs:
        title = kwargs["title"]
    else:
        return "输入参数错误，根据工具参数要求输入title参数"

    desc = ""
    if "desc" in kwargs:
        desc = kwargs["desc"]
    else:
        return "输入参数错误，根据工具参数要求输入desc参数"

    data = extract_data(desc)

    res = {
        "title": title,
        "type": "barChart",
        "data": json.loads(data),
    }

    web_json.append(res)

    # return json.dumps(res, ensure_ascii=False)
    return "柱状图生成成功"
